/*
 * CudaAesOperationsTest.cpp
 *
 *  Created on: Feb 27, 2010
 *      Author: Łukasz Spintzyk
 */
//#define DBG_NO_CUDA
#ifndef DBG_NO_CUDA
#include "CudaAESInvOperations.h"
#include "CudaAESOperations.h"
#endif
#include "aesUtils.h"

#include <iostream>
#include <iomanip>

#include <cppunit/TestRunner.h>
#include <cppunit/TestResult.h>
#include <cppunit/TestResultCollector.h>
#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/BriefTestProgressListener.h>
#include <cppunit/extensions/TestFactoryRegistry.h>

class TestInvRoundTransformations: public CPPUNIT_NS::TestCase
{
CPPUNIT_TEST_SUITE(TestInvRoundTransformations);
		CPPUNIT_TEST(testInvSubByte4x4);
		CPPUNIT_TEST(testInvSubByte4x6);
		CPPUNIT_TEST(testInvSubByte4x8);
		CPPUNIT_TEST(testInvShiftRow4x4);
		CPPUNIT_TEST(testInvShiftRow4x6);
		CPPUNIT_TEST(testInvShiftRow4x8);
		CPPUNIT_TEST(testInvMixColumn4x4);
		CPPUNIT_TEST(testInvMixColumn4x6);
		CPPUNIT_TEST(testInvMixColumn4x8);

	CPPUNIT_TEST_SUITE_END();

public:
	void setUp(void)
	{
		d_blockSize = 4 * 8;
		d_block = new byte[d_blockSize];
		d_copyBlock = new byte[d_blockSize];
		d_copyBlock2 = new byte[d_blockSize];
		d_roundKey = new byte[d_blockSize];

		setTestBlock(d_block, 8, 4);
		setTestBlock(d_roundKey, 8, 4, 152);
		memcpy(d_copyBlock, d_block, sizeof(byte) * d_blockSize);
		memcpy(d_copyBlock2, d_block, sizeof(byte) * d_blockSize);
		aesDecryptInit();
	}
	void tearDown(void)
	{
		for (int i = 0; i < d_blockSize; ++i)
		{
			d_block[i] = d_copyBlock[i] = 0;
		}

		delete[] d_block;
		delete[] d_copyBlock;
		delete[] d_copyBlock2;
		delete[] d_roundKey;
	}

protected:
	void testInvSubByte4x4(void)
	{
		const unsigned int size = 4 * 4;
		invByteSub(d_block, NB_4);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX_INV[d_copyBlock[i]]);
		}
	}

	void testInvSubByte4x6(void)
	{
		const unsigned int size = 4 * 6;
		invByteSub(d_block, NB_6);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX_INV[d_copyBlock[i]]);
		}
	}

	void testInvSubByte4x8(void)
	{
		const unsigned int size = 4 * 8;
		invByteSub(d_block, NB_8);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX_INV[d_copyBlock[i]]);
		}
	}

	void testInvShiftRow4x4(void)
	{
		const unsigned int size = 4 * 4;
		invShiftRow(d_block, NB_4);
		invShiftRow2(d_copyBlock, NB_4);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);

		shiftRow(d_block, NB_4);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock2, size) == 0);
	}

	void testInvShiftRow4x6(void)
	{
		const unsigned int size = 4 * 6;
		invShiftRow(d_block, NB_6);
		invShiftRow2(d_copyBlock, NB_6);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);

		shiftRow(d_block, NB_6);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock2, size) == 0);
	}

	void testInvShiftRow4x8(void)
	{
		const int size = 4 * 8;

		invShiftRow(d_block, NB_8);
		invShiftRow2(d_copyBlock, NB_8);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);

		shiftRow(d_block, NB_8);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock2, size) == 0);
	}

	void testInvMixColumn4x4(void)
	{
		const int size = 4 * 4;

		d_copyBlock[4+0] = d_block[4+0] = 0x04;
		d_copyBlock[4+1] = d_block[4+1] = 0x66;
		d_copyBlock[4+2] = d_block[4+2] = 0x81;
		d_copyBlock[4+3] = d_block[4+3] = 0xe5;


		invMixColumn(d_block, NB_4);
		invMixColumnTest(d_copyBlock, 4, 4, (byte*) INVMIXTAB, 4, 4);


		CPPUNIT_ASSERT(d_block[4+0]==0xd4);
		CPPUNIT_ASSERT(d_block[4+1]==0xbf);
		CPPUNIT_ASSERT(d_block[4+2]==0x5d);
		CPPUNIT_ASSERT(d_block[4+3]==0x30);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);

	}
	void testInvMixColumn4x6(void)
	{
		const int size = 4 * 6;

		invMixColumn(d_block, NB_6);
		invMixColumnTest(d_copyBlock, 6, 4, (byte*) INVMIXTAB, 4, 4);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}
	void testInvMixColumn4x8(void)
	{

		const int size = 4 * 8;

		invMixColumn(d_block, NB_8);
		invMixColumnTest(d_copyBlock, 8, 4, (byte*) INVMIXTAB, 4, 4);


		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}

private:
	void invShiftRow2(byte* state, NB nb)
	{
		unsigned int size = BLOCK_SIZE(nb);
		byte* state_copy = new byte[size * sizeof(byte)];
		memcpy(state_copy, state, size);

		int w = size / 4;
		for (int x = 0; x < w; ++x)
		{
			int x2 = (x + C1(nb)) % w;
			state_copy[x2 * 4 + 1] = state[x * 4 + 1];
			assert(x2*4+1 < size);
			assert(x*4+1 < size);

			x2 = (x + C2(nb)) % w;
			state_copy[x2 * 4 + 2] = state[x * 4 + 2];
			x2 = (x + C3(nb)) % w;
			state_copy[x2 * 4 + 3] = state[x * 4 + 3];
		}
		memcpy(state, state_copy, size);
		delete[] state_copy;
	}

	void setTestBlock(byte* block, int w, int h, int seed = 1)
	{
		int k = seed;
		for (int i = 0; i < w; ++i)
		{
			for (int j = 0; j < h; ++j)
			{
				block[i * h + j] = k++;
			}
		}
	}

	void printBlock(byte* block, int w, int h)
	{
		for (int j = 0; j < h; ++j)
		{
			for (int i = 0; i < w; ++i)
			{
				std::cout << std::hex << std::setw(4) << (unsigned int) block[i
						* h + j] << " ";
			}
			std::cout << std::endl;
		}
	}

	void invMixColumnTest(byte* state, int w1, int h1,
			byte* CBlock, int w2, int h2)
	{

		for (int i = 0; i < w1; ++i)
		{
			byte* column = new byte[h1];
			for (int j = 0; j < h1; ++j)
			{
				byte* p = state + i * h1 + j;
				column[j] = *p;
			}
			assert(h1==4);
			assert(w2 == 4 );
			assert(h2 == 4);
			matMul(column, h1, CBlock, w2, h2);
			for (int j = 0; j < h1; ++j)
			{
				byte* p = state + i * h1 + j;
				*p = column[j];
			}
			delete[] column;
		}
	}

	byte field28Mno(byte c1, byte c2)
	{
		byte result = 0;

		int count = 8;
		while (count--)
		{
			if (c2 & 0x01)
				result ^= c1;
			if (c1 & 0x80)
			{
				c1 <<= 1;
				c1 ^= 0x1B;
			}
			else
				c1 <<= 1;
			c2 >>= 1;
		}
		return result;

	}

	void matMul(byte* a, int ha, byte* b, int wb, int hb)
	{
		byte* copy = new byte[ha];

		assert(ha == hb);

		for (int i = 0; i < ha; ++i)
		{
			byte sum = 0;
			sum = 0;
			for (int j = 0; j < wb; ++j)
			{
				byte* p = b + i * ha + j;
				sum ^= field28Mno(a[j], *p);//*p * a[j];
			}
			copy[i] = sum;
		}

		memcpy(a, copy, ha);

		delete[] copy;
	}

	byte* d_block;
	byte* d_copyBlock;
	byte* d_copyBlock2;
	byte* d_roundKey;
	int d_blockSize;

};
CPPUNIT_TEST_SUITE_REGISTRATION(TestInvRoundTransformations);

int main(int ac, char **av)
{
	//--- Create the event manager and test controller
	CPPUNIT_NS::TestResult controller;
	//--- Add a listener that colllects test result
	CPPUNIT_NS::TestResultCollector result;
	controller.addListener(&result);
	//--- Add a listener that print dots as test run.
	CPPUNIT_NS::BriefTestProgressListener progress;
	controller.addListener(&progress);
	//--- Add the top suite to the test runner
	CPPUNIT_NS::TestRunner runner;
	runner.addTest(CPPUNIT_NS::TestFactoryRegistry::getRegistry().makeTest());
	runner.run(controller);

	return result.wasSuccessful() ? 0 : 1;
}
